import torch
from torch import nn
from dataset import myDataSet
from torch.utils.data import DataLoader
from u2net import U2NETP, U2NET
import os
import numpy as np
from torchvision.utils import save_image
import cfg
# from tools import SegmentationMetric, mean_iou
from loss import OhemCELoss, SoftDiceLoss
from PIL import Image
import cv2
import math
from test_img import Run

bce_loss = nn.BCELoss(size_average=True)
dice_loss = SoftDiceLoss()


# bce_loss = OhemCELoss(thresh=0.1)


def muti_bce_loss_fusion(d0, d1, d2, d3, d4, d5, d6, labels_v,i):
    loss0 = bce_loss(d0, labels_v)
    loss1 = bce_loss(d1, labels_v)
    loss2 = bce_loss(d2, labels_v)
    loss3 = bce_loss(d3, labels_v)
    loss4 = bce_loss(d4, labels_v)
    loss5 = bce_loss(d5, labels_v)
    loss6 = bce_loss(d6, labels_v)

    # loss0_dice = dice_loss(d0, labels_v)
    # loss1_dice = dice_loss(d1, labels_v)
    # loss2_dice = dice_loss(d2, labels_v)
    # loss3_dice = dice_loss(d3, labels_v)
    # loss4_dice = dice_loss(d4, labels_v)
    # loss5_dice = dice_loss(d5, labels_v)
    # loss6_dice = dice_loss(d6, labels_v)

    loss = loss0 + loss1 + loss2 + loss3 + loss4 + loss5 + loss6
    # loss = loss + loss0_dice + loss1_dice + loss2_dice + loss3_dice + loss4_dice + loss5_dice + loss6_dice
    if (i%50==0):print("l0: %3f, l1: %3f, l2: %3f, l3: %3f, l4: %3f, l5: %3f, l6: %3f\n" % (
        loss0.item(), loss1.item(), loss2.item(), loss3.item(), loss4.item(), loss5.item(), loss6.item()))

    return loss0, loss


def dice_coeff(pred, target):
    smooth = 1.
    num = pred.size(0)
    m1 = pred.view(num, -1)  # Flatten
    m2 = target.view(num, -1)  # Flatten
    intersection = (m1 * m2).sum()  # m1与m2的交集

    return (2. * intersection + smooth) / (m1.sum() + m2.sum() + smooth)

def tensor_img(tens):
    img=np.array(tens,dtype=np.float32)
    return Image.fromarray(img)

# 计算psnr
def psnr(img1, img2):
   mse = np.mean((img1/1.0 - img2/1.0) ** 2 )
   if mse < 1.0e-10:
      return 100
   return 10 * math.log10(255.0**2/mse)


if __name__ == '__main__':

    start_epoch = 188
    netSavePath = 'params/168_.83.797126.pth'
    resultPath = r'./data/result_imgs'

    trainData = myDataSet('data/train0_crop', is_train=True)
    train_dataLoader = DataLoader(trainData, batch_size=cfg.batch_size, shuffle=True)

    testData = myDataSet('data/test0_crop', is_train=False)
    test_dataLoader = DataLoader(testData, batch_size=1, shuffle=True)

    # 设定程序运行设备
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

    # 设定运行网络
    net = U2NET().to(device)

    # 加载参数
    if os.path.exists(netSavePath):
        net.load_state_dict(torch.load(netSavePath))
        print('u2net神经网络参数加载成功！')
    else:
        print('No Params')

    opt = torch.optim.Adam(net.parameters())

    net.train()

    loss_list = []
    best_psnr = 0
    for episode in range(100000):
        start_epoch += 1
        for i, (dataImg, labelImg,_,_) in enumerate(train_dataLoader):
            if(i>500):break
            dataImg = dataImg.to(device).float()
            labelImg=labelImg.to(device)[:,:1].float()

            # 计算loss损失
            d0, d1, d2, d3, d4, d5, d6 = net(dataImg)

            loss2, loss = muti_bce_loss_fusion(d0, d1, d2, d3, d4, d5, d6, labelImg.float(),i)
            loss_list.append(loss.item())

            if i % 6 == 0: opt.zero_grad()
            loss.backward()
            if i % 6 == 0: opt.step()

        test_dice = []


        test_psnr = []
        # run=Run('params/8_.80.337132.pth')
        with torch.no_grad():
            net.eval()
            for i, (dataImg, labelImg,orign_img,img_index) in enumerate(test_dataLoader):
                dataImg = dataImg.to(device).float()
                labelImg = labelImg.to(device).float()
                # clip_mask=run.generate_mask(dataImg)
                

                # 计算loss损失
                d0, d1, d2, d3, d4, d5, d6 = net(dataImg)

                d0 = d0.cpu().detach()
                d0 = torch.where(d0 > 0.5, 1., 0.)

                if(0):
                    h,w=labelImg[0,0].shape
                    scale_rate=0.7
                    n_h=int(scale_rate*h)
                    n_w=int(scale_rate*w)
                    new_img=np.hstack((cv2.resize(d0.numpy()[0, 0]*255,(n_w,n_h)),cv2.resize(labelImg.cpu().numpy()[0,0]*255,(n_w,n_h))))
                    cv2.imshow("img",new_img)
                    # cv2.imshow("img",labelImg.cpu().numpy()[0,0]*255.)
                    cv2.waitKey(0)
                if(0):
                    img=Image.open(os.path.join(testData.root_path,"images",testData.file_name_list[img_index[0]]))
                    _img  = np.array(img)
                    if (len(_img.shape)==2):_img=np.repeat(_img[...,None],3,axis=2)

                    target_mask=Image.open(os.path.join(testData.root_path,"mask",testData.file_name_list[img_index[0]]))

                    inputImage=testData.img2tensor(Image.fromarray(_img))
                    inputImage=inputImage[None,...]

                    _,_,h,w,=inputImage.shape
                    rh, rw = h, w
                    step = cfg.img_size
                    pad_h = step - h if h < step else 0
                    pad_w = step - w if w < step else 0
                    m =torch.nn.ConstantPad2d((0, pad_w, 0, pad_h),0.)
                    imgs=m(inputImage)

                    _, _, h, w = imgs.shape
                    res = np.zeros_like(imgs)

                    kernel=cv2.getStructuringElement(cv2.MORPH_RECT,(6,6))

                    for i in range(0, h, step):
                        for j in range(0, w, step):
                            if h - i < step:
                                i = h - step
                            if w - j < step:
                                j = w - step
                            clip = imgs[:, :, i:i + step, j:j + step]
                            clip = clip.cuda()
                            d0, d1, d2, d3, d4, d5, d6 = net(clip)
                            d0 = torch.where(d0 > 0.5, 255., 0.)
                            out_mask = d0[0,0].cpu().numpy()
                            out_mask=np.repeat(out_mask[...,None],3,axis=2)
                            
                            
                            clip = clip.cpu().numpy()
                            
                            if(0):
                                dilate=cv2.morphologyEx(out_mask,cv2.MORPH_DILATE,kernel,iterations=10)

                            if(0):
                                out_clip=np.transpose(clip[0],(1,2,0))[...,::-1]
                                o=np.hstack((out_mask,out_clip))
                                cv2.imshow("mask",o)
                                cv2.waitKey(0)

                            out_mask=out_mask[...,0]
                            res[:,:, i:i + step, j:j + step] = out_mask
                            del out_mask, clip
                    
                    res = res[:,:, :rh, :rw]
                    res=np.array(np.transpose(res[0],(1,2,0)),dtype=np.float32)
                    cv2.imwrite("img.jpg",res)
                    exit()
                
                psnr_value=psnr(d0[0, 0].float().numpy(),labelImg.cpu()[0,0].numpy())

                if(0 and i%20==0):
                    out_mask=np.array(d0[0, 0]*255.)
                    target_mask=np.array(labelImg.cpu()[0, 0]*255.)
                    new_img=np.hstack((out_mask,target_mask))
                   
                    if(0):
                        cv2.imshow("scr", new_img)
                        cv2.waitKey(0)
                    cv2.imwrite(f"data/result_imgs/{start_epoch}_{i}.jpg",new_img)

                test_psnr.append(psnr_value)
                if start_epoch>1 and i%50==0:print(f"epoch:{start_epoch},psnr:{psnr_value},best_psnr:{best_psnr}")

        # if start_epoch%2==0:
        torch.save(net.state_dict(),'params/{}_.{:3f}.pth'. format(start_epoch, np.mean(test_psnr)))
        avg_psnr=np.mean(test_psnr)
        if (avg_psnr > best_psnr):
            best_psnr = np.mean(test_psnr)
            torch.save(net.state_dict(), 'params.pth')
            print('Model parameters saved successfully !')
